# coding: utf-8

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
from keras import backend as K
from keras.layers.convolutional import Conv2D
from keras.layers.core import Dense, Flatten
from keras.layers.pooling import MaxPooling2D
from tensorflow.examples.tutorials.mnist import input_data

# Using TensorFlow backend.
print(K.image_data_format())

data_dir = '/tmp/tensorflow/mnist/input_data'
mnist = input_data.read_data_sets(data_dir, one_hot=True)

# Create the model
x = tf.placeholder(tf.float32, [None, 784])
y_ = tf.placeholder(tf.float32, [None, 10])
learning_rate = tf.placeholder(tf.float32)

with tf.name_scope("reshape"):
    x_image = tf.reshape(x, [-1, 28, 28, 1])

# 输出32个通过， 换个理解就是使用了 32 个 5*5*1 的卷积核进行卷积操作
net1 = Conv2D(32, kernel_size=[5, 5], strides=[1, 1], activation="relu", padding="same", input_shape=[28, 28, 1])(
    x_image)

# 池化，维度不变化
net2 = MaxPooling2D(pool_size=[2, 2])(net1)

net3 = Conv2D(64, kernel_size=[5, 5], strides=[1, 1], activation="relu", padding="same")(net2)

net4 = MaxPooling2D(pool_size=[2, 2])(net3)

net5 = Flatten()(net4)

net6 = Dense(1000, activation="relu")(net5)

net7 = Dense(10, activation="softmax")(net6)

from keras.objectives import categorical_crossentropy
cross_entropy = tf.reduce_mean(categorical_crossentropy(y_, net7))

l2_loss = tf.add_n( [tf.nn.l2_loss(w) for w in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)] )
total_loss = cross_entropy + 7e-5*l2_loss

train_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(total_loss)

sess = tf.Session()

K.set_session(sess)

init_op = tf.global_variables_initializer()
sess.run(init_op)
# Train
for step in range(3000):
    batch_xs, batch_ys = mnist.train.next_batch(100)
    lr = 0.01
    _, loss, l2_loss_value, total_loss_value = sess.run(
        [train_step, cross_entropy, l2_loss, total_loss],
        feed_dict={x: batch_xs, y_: batch_ys, learning_rate: lr})

    if (step + 1) % 100 == 0:
        print('step %d, entropy loss: %f, l2_loss: %f, total loss: %f' %
              (step + 1, loss, l2_loss_value, total_loss_value))
        # Test trained model
        correct_prediction = tf.equal(tf.argmax(net7, 1), tf.argmax(y_, 1))
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
        print(sess.run(accuracy, feed_dict={x: batch_xs, y_: batch_ys}))
    if (step + 1) % 1000 == 0:
        print(sess.run(accuracy, feed_dict={x: mnist.test.images,
                                            y_: mnist.test.labels}))
